#!/usr/bin/env python3
import os
import sys
import subprocess
import readline
import shlex
import json
from typing import List, Optional
from openai import OpenAI

# ==================== AI 模块部分 ====================

# 全局变量存储对话记忆
_conversation_memory = []
_MAX_MEMORY_ITEMS = 20  # 最大记忆条数

def fix_json_format(content):
    """
    修复可能的JSON格式问题，特别是处理多余的括号
    
    参数:
        content: 原始响应内容字符串
    
    返回:
        修复后的内容字符串
    """
    # 去除首尾空白字符
    content = content.strip()
    
    # 检查是否以{开头且{比}多一个
    if content.startswith('{'):
        open_count = content.count('{')
        close_count = content.count('}')
        
        # 如果{比}多一个，尝试去掉开头的{
        if open_count == close_count + 1:
            content = content[1:].strip()
            print("修复了JSON格式: 移除了多余的开头'{'")
    
    return content

def clear_memory():
    """
    清空对话记忆
    """
    global _conversation_memory
    _conversation_memory = []
    print("对话记忆已清空")

def get_conversation_history():
    """
    获取当前对话历史
    
    返回:
        对话历史列表
    """
    return _conversation_memory.copy()

def set_memory_limit(limit):
    """
    设置记忆上限
    
    参数:
        limit: 最大记忆条数
    """
    global _MAX_MEMORY_ITEMS
    _MAX_MEMORY_ITEMS = max(1, limit)  # 至少保留1条
    print(f"记忆上限已设置为 {_MAX_MEMORY_ITEMS} 条")

def get_memory_limit():
    """
    获取当前记忆上限
    
    返回:
        最大记忆条数
    """
    return _MAX_MEMORY_ITEMS

def _add_to_memory(user_message, assistant_response, user_feedback=None, command_result=None):
    """
    内部函数：添加对话到记忆，自动管理上限
    
    参数:
        user_message: 用户消息内容
        assistant_response: 助手响应内容
        user_feedback: 用户反馈（如拒绝执行等）
        command_result: 命令执行结果（可选）
    """
    global _conversation_memory
    
    # 添加用户消息
    _conversation_memory.append({"role": "user", "content": user_message})
    
    # 构建助手响应内容
    assistant_content = assistant_response
    
    # 添加用户反馈（如拒绝执行）
    if user_feedback:
        assistant_content += f"\n\n用户反馈: {user_feedback}"
    
    # 添加命令执行结果
    if command_result:
        assistant_content += f"\n\n命令执行结果:\n{command_result}"
    
    # 添加助手响应
    _conversation_memory.append({"role": "assistant", "content": assistant_content})
    
    # 检查并限制记忆长度
    while len(_conversation_memory) > _MAX_MEMORY_ITEMS * 2:  # 乘以2因为每条对话包含user和assistant
        # 移除最旧的一对对话（user和assistant）
        if len(_conversation_memory) >= 2:
            _conversation_memory = _conversation_memory[2:]
    
    print(f"对话记忆已更新，当前共有 {len(_conversation_memory)//2} 轮对话（上限: {_MAX_MEMORY_ITEMS} 轮）")

def update_command_result(command_result, use_memory=True):
    """
    更新最近一条命令的执行结果到记忆中
    
    参数:
        command_result: 命令执行结果
        use_memory: 是否使用记忆功能
    """
    if not use_memory or not _conversation_memory:
        return
    
    # 找到最近的一条助手消息
    for i in range(len(_conversation_memory)-1, -1, -1):
        if _conversation_memory[i]["role"] == "assistant":
            # 检查是否已经包含命令结果
            current_content = _conversation_memory[i]["content"]
            
            # 如果已经包含用户反馈，在用户反馈后添加命令结果
            if "用户反馈:" in current_content and "命令执行结果:" not in current_content:
                # 在用户反馈后添加命令结果
                parts = current_content.split("用户反馈:")
                if len(parts) >= 2:
                    feedback_part = parts[1]
                    # 检查是否有其他内容在用户反馈之后
                    if "\n\n" in feedback_part:
                        feedback_parts = feedback_part.split("\n\n", 1)
                        new_content = parts[0] + f"用户反馈:{feedback_parts[0]}\n\n命令执行结果:\n{command_result}"
                        if len(feedback_parts) > 1:
                            new_content += f"\n\n{feedback_parts[1]}"
                    else:
                        new_content = parts[0] + f"用户反馈:{feedback_part}\n\n命令执行结果:\n{command_result}"
                    _conversation_memory[i]["content"] = new_content
                    print("命令执行结果已添加到记忆（在用户反馈后）")
            elif "命令执行结果:" in current_content:
                # 如果已经存在命令结果，更新它
                parts = current_content.split("命令执行结果:")
                if len(parts) >= 2:
                    # 保留第一部分，更新结果部分
                    result_part = parts[1]
                    # 检查是否有其他内容在命令结果之后
                    if "\n\n" in result_part:
                        result_parts = result_part.split("\n\n", 1)
                        new_content = parts[0] + f"命令执行结果:\n{command_result}"
                        if len(result_parts) > 1:
                            new_content += f"\n\n{result_parts[1]}"
                    else:
                        new_content = parts[0] + f"命令执行结果:\n{command_result}"
                    _conversation_memory[i]["content"] = new_content
                    print("命令执行结果已更新")
            else:
                # 直接添加命令结果
                _conversation_memory[i]["content"] += f"\n\n命令执行结果:\n{command_result}"
                print("命令执行结果已添加到记忆")
            break

def get_command(api_key, system_prompt, user_prompt, model_name, enable_thinking=False, use_memory=True, command_result=None):
    """
    调用 OpenAI 风格 API 获取命令
    
    参数:
        api_key: API 密钥
        system_prompt: 系统提示词
        user_prompt: 用户提示词
        model_name: 模型名称
        enable_thinking: 是否启用深度思考功能
        use_memory: 是否使用对话记忆功能
        command_result: 上一条命令的执行结果（用于添加到记忆）
    
    返回:
        命令字符串或 None
    """
    global _conversation_memory
    
    try:
        # 创建客户端
        client = OpenAI(
            api_key=api_key,
            base_url="https://aistudio.baidu.com/llm/lmapi/v3"
        )

        # 准备消息 - 强制要求 JSON 格式响应
        full_system_prompt = f"""{system_prompt}

你必须以严格的 JSON 格式响应，且只包含一个 'cmd' 字段，其值为要在终端中执行的命令。
例如：{{"cmd": "ls -la"}}
不要添加任何额外格式，如代码块标记或解释文本。"""

        # 构建消息列表
        messages = [{"role": "system", "content": full_system_prompt}]
        
        # 如果启用记忆功能且存在历史对话，添加历史记录
        if use_memory and _conversation_memory:
            # 如果有上一条命令的执行结果，需要先更新记忆
            if command_result and len(_conversation_memory) >= 2:
                # 找到最近的一条助手消息，添加命令结果
                latest_assistant_index = None
                for i in range(len(_conversation_memory)-1, -1, -1):
                    if _conversation_memory[i]["role"] == "assistant":
                        latest_assistant_index = i
                        break
                
                if latest_assistant_index is not None:
                    # 检查是否已经包含命令结果
                    current_content = _conversation_memory[latest_assistant_index]["content"]
                    if "命令执行结果:" not in current_content and "用户反馈:" not in current_content:
                        _conversation_memory[latest_assistant_index]["content"] += f"\n\n命令执行结果:\n{command_result}"
                    elif "用户反馈:" in current_content and "命令执行结果:" not in current_content:
                        # 在用户反馈后添加命令结果
                        parts = current_content.split("用户反馈:")
                        if len(parts) >= 2:
                            feedback_part = parts[1]
                            if "\n\n" in feedback_part:
                                feedback_parts = feedback_part.split("\n\n", 1)
                                new_content = parts[0] + f"用户反馈:{feedback_parts[0]}\n\n命令执行结果:\n{command_result}"
                                if len(feedback_parts) > 1:
                                    new_content += f"\n\n{feedback_parts[1]}"
                            else:
                                new_content = parts[0] + f"用户反馈:{feedback_part}\n\n命令执行结果:\n{command_result}"
                            _conversation_memory[latest_assistant_index]["content"] = new_content
            
            # 添加历史对话
            messages.extend(_conversation_memory)
            print(f"已加载 {len(_conversation_memory)//2} 轮历史对话")
        
        # 添加当前用户提示
        messages.append({"role": "user", "content": user_prompt})

        # 准备请求参数
        request_params = {
            "model": model_name,
            "messages": messages,
            "temperature": 0.1,
            "response_format": {"type": "json_object"},  # 强制要求 JSON 响应
            "stream": True  # 启用流式响应
        }
        
        # 如果启用深度思考，添加metadata参数（需要将布尔值转换为字符串）
        # FIXME: 这个暂时没有用
        if enable_thinking:
            request_params["metadata"] = {"enable_thinking": 'true'}  # 改为字符串

        # 调用 API - 使用流式响应以获取思维链
        completion = client.chat.completions.create(**request_params)
        
        # 收集完整响应内容
        full_content = ""
        reasoning_content = ""
        has_reasoning = False
        has_content = False
        
        # 灰色 ANSI 颜色代码
        GRAY_COLOR = "\x1b[30m"
        RESET_COLOR = "\x1b[0m"
        
        for chunk in completion:
            if len(chunk.choices) > 0:
                # 检查是否有思维链内容
                if hasattr(chunk.choices[0].delta, 'reasoning_content') and chunk.choices[0].delta.reasoning_content:
                    reasoning_chunk = chunk.choices[0].delta.reasoning_content
                    reasoning_content += reasoning_chunk
                    # 打印灰色思维链
                    print(f"{GRAY_COLOR}{reasoning_chunk}{RESET_COLOR}", end="", flush=True)
                    has_reasoning = True
                
                # 检查是否有普通内容
                if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content:
                    content_chunk = chunk.choices[0].delta.content
                    full_content += content_chunk
                    # 只有在没有思维链时才打印普通内容
                    #if not has_reasoning:
                    #    print(content_chunk, end="", flush=True)
                    #    has_content = True
        
        # 如果有思维链内容，打印换行符
        if has_reasoning:
            print("\n")  # 在思维链后添加回车换行
        
        # 解析最终响应内容
        command = None
        if full_content:
            # 修复可能的JSON格式问题
            fixed_content = fix_json_format(full_content)
            
            # 尝试解析修复后的JSON
            try:
                data = json.loads(fixed_content)
                if "cmd" in data:
                    command = data["cmd"]
                else:
                    print(f"响应中缺少 'cmd' 字段: {fixed_content}")
            except json.JSONDecodeError as e:
                print(f"修复后仍无法解析JSON: {e}, 原始内容: {full_content}, 修复后内容: {fixed_content}")
        
        return command, full_content
        
    except Exception as e:
        print(f"API 调用错误: {e}")
        return None, None

# ==================== Shell 部分 ====================

class AIShell:
    def __init__(self):
        self.history_file = os.path.expanduser("~/.aish_history")
        self.setup_readline()
        
        # 从环境变量读取AI配置
        self.api_key = os.getenv('BAIDU_FEIJIANG_API_KEY', '').strip()
        self.system_prompt = os.getenv('WENXIN_SYSTEM_PROMPT', "")
        
        # 记忆功能设置
        self.use_memory = True  # 默认启用记忆功能
        self.last_command_result = None  # 存储上一条命令的执行结果
        
        # 检查必要的环境变量
        if not self.api_key:
            print("警告: 未设置 BAIDU_FEIJIANG_API_KEY 环境变量")
            print("请设置环境变量或通过启动脚本设置")
        
    def setup_readline(self):
        """设置 readline 支持历史记录"""
        try:
            readline.read_history_file(self.history_file)
        except FileNotFoundError:
            pass
        readline.set_history_length(1000)
        
    def save_history(self):
        """保存历史记录"""
        try:
            readline.write_history_file(self.history_file)
        except Exception:
            pass

    def ask_wenxin(self, query: str, model: str, enable_thinking: bool) -> Optional[str]:
        """
        调用AI模块获取命令建议
        """
        if not query.strip():
            print("错误: 查询内容为空")
            return None
            
        # 检查API密钥
        if not self.api_key:
            print("错误: 未设置百度飞桨API密钥")
            print("请设置 BAIDU_FEIJIANG_API_KEY 环境变量")
            return None
            
        print(f"正在向AI模型 {model} 请求建议...")
        
        # 调用AI模块（传入上一条命令的执行结果）
        suggested_command, full_content = get_command(
            api_key=self.api_key,
            system_prompt=self.system_prompt,
            user_prompt=query,
            model_name=model,
            enable_thinking=enable_thinking,
            use_memory=self.use_memory,
            command_result=self.last_command_result
        )
        
        return suggested_command, full_content

    def print_help(self):
        """显示帮助信息"""
        print("aish - AI Shell")
        print()
        print("使用说明:")
        print("  直接输入常规 shell 命令执行本地操作")
        print("  以 ? 或 ?? 开头输入自然语言指令调用AI大模型")
        print("  支持对话记忆功能，AI会记住之前的对话和命令执行结果")
        print("  包括用户拒绝执行命令的反馈也会被记录")
        print()
        print("命令格式:")
        print("  command [args]              # 执行常规shell命令")
        print("  ? [自然语言指令]            # 调用AI模型（不启用思考模式）")
        print("  ?? [自然语言指令]           # 调用AI模型（启用思考模式）")
        print("  memory                      # 显示记忆状态")
        print("  memory clear                # 清空记忆")
        print("  memory limit [数字]         # 设置记忆上限")
        print("  memory on/off               # 启用/禁用记忆功能")
        print("  exit                        # 退出shell")
        print("  help, --help, -h            # 显示此帮助信息")

    def handle_memory_command(self, command_parts: List[str]):
        """处理记忆相关命令"""
        if len(command_parts) == 1:
            # 显示记忆状态
            current_memory = len(get_conversation_history()) // 2
            memory_limit = get_memory_limit()
            status = "启用" if self.use_memory else "禁用"
            print(f"记忆功能: {status}")
            print(f"当前记忆: {current_memory} 轮对话")
            print(f"记忆上限: {memory_limit} 轮对话")
            
            # 显示最近几条记忆内容
            if current_memory > 0:
                print("\n最近对话:")
                history = get_conversation_history()
                # 显示最近3轮对话
                start_idx = max(0, len(history) - 6)  # 每轮对话包含2条消息
                for i in range(start_idx, len(history), 2):
                    if i + 1 < len(history):
                        user_msg = history[i]['content'][:50] + "..." if len(history[i]['content']) > 50 else history[i]['content']
                        assistant_msg = history[i+1]['content'][:50] + "..." if len(history[i+1]['content']) > 50 else history[i+1]['content']
                        print(f"  {i//2 + 1}. 用户: {user_msg}")
                        print(f"     助手: {assistant_msg}")
            return True
        elif len(command_parts) == 2:
            if command_parts[1] == "clear":
                clear_memory()
                self.last_command_result = None
                return True
            elif command_parts[1] == "on":
                self.use_memory = True
                print("记忆功能已启用")
                return True
            elif command_parts[1] == "off":
                self.use_memory = False
                print("记忆功能已禁用")
                return True
        elif len(command_parts) == 3 and command_parts[1] == "limit":
            try:
                limit = int(command_parts[2])
                set_memory_limit(limit)
                return True
            except ValueError:
                print("错误: 请输入有效的数字")
                return True
        return False

    def get_prompt(self) -> str:
        """获取提示符"""
        current_dir = os.getcwd()
        username = os.getenv('USER', 'unknown')
        hostname = os.getenv('HOSTNAME', 'unknown')
        
        # 在提示符中显示记忆状态
        memory_indicator = "[M]" if self.use_memory else "[X]"
        
        return f"{memory_indicator} {username}@{hostname}:{current_dir}$ "

    def execute_command(self, command: List[str]) -> int:
        """执行命令并记录结果"""
        try:
            # 特殊处理 cd 命令
            if command[0] == "cd":
                if len(command) > 1:
                    new_dir = command[1]
                    os.chdir(new_dir)
                    result_msg = f"目录切换成功 -> {os.getcwd()}"
                else:
                    home_dir = os.path.expanduser("~")
                    os.chdir(home_dir)
                    result_msg = f"目录切换成功 -> {home_dir}"
                
                self.last_command_result = result_msg
                print(result_msg)
                return 0
                
            # 执行其他命令
            result = subprocess.run(command, shell=False, capture_output=True, text=True)
            
            # 记录命令执行结果
            if result.stdout:
                self.last_command_result = result.stdout
            elif result.stderr:
                self.last_command_result = result.stderr
            else:
                self.last_command_result = f"命令执行完成，退出码: {result.returncode}"
            
            # 输出结果
            if result.stdout:
                print(result.stdout)
            if result.stderr:
                print(result.stderr, file=sys.stderr)
                
            # 如果启用了记忆功能，更新命令结果到记忆
            if self.use_memory:
                update_command_result(self.last_command_result, self.use_memory)
                
            return result.returncode
        except Exception as e:
            error_msg = f"执行错误: {e}"
            print(error_msg)
            self.last_command_result = error_msg
            return 1

    def handle_ai_command(self, input_str: str):
        """处理以 ? 或 ?? 开头的 AI 命令"""
        if input_str.startswith("??"):
            # 双问号命令，启用思考模式
            query = input_str[2:].strip()
            model = "ernie-x1-turbo-32k"
            enable_thinking = True
        elif input_str.startswith("?"):
            # 单问号命令，不启用思考模式
            query = input_str[1:].strip()
            model = "ernie-4.5-turbo-32k"
            enable_thinking = False
        else:
            return
            
        if not query:
            print("错误: 请在 ? 或 ?? 后输入查询内容")
            return
            
        # 调用AI获取建议
        suggested_command, full_content = self.ask_wenxin(query, model, enable_thinking)
        
        if suggested_command:
            print(f"\nAI建议执行: {suggested_command}")
            
            # 询问用户是否执行
            while True:
                response = input("是否执行此命令? (y/N): ").strip().lower()
                if response in ['y', 'yes']:
                    # 执行建议的命令
                    command = self.parse_command(suggested_command)
                    if command:
                        self.execute_command(command)
                    
                    # 更新记忆（包含命令执行结果）
                    if self.use_memory and full_content:
                        _add_to_memory(query, full_content, command_result=self.last_command_result)
                    break
                elif response in ['', 'n', 'no']:
                    user_feedback = "用户拒绝执行该命令"
                    print(user_feedback)
                    
                    # 更新记忆（包含用户拒绝执行的反馈）
                    if self.use_memory and full_content:
                        _add_to_memory(query, full_content, user_feedback=user_feedback)
                    break
                else:
                    print("请输入 y 或 n")
        else:
            error_msg = "未能获取到有效的命令建议"
            print(error_msg)
            # 即使获取建议失败，也记录到记忆中
            if self.use_memory:
                _add_to_memory(query, error_msg)

    def parse_command(self, input_str: str) -> List[str]:
        """解析命令字符串"""
        try:
            return shlex.split(input_str)
        except ValueError as e:
            print(f"解析错误: {e}")
            return []

    def run(self):
        """主循环"""
        print("aish - AI Shell 已启动。输入 'help' 查看帮助，'exit' 退出。")
        print(f"记忆功能已{'启用' if self.use_memory else '禁用'}，上限: {get_memory_limit()} 轮对话")
        
        while True:
            try:
                # 获取输入
                prompt = self.get_prompt()
                try:
                    user_input = input(prompt).strip()
                except EOFError:
                    print()  # 换行后退出
                    break
                except KeyboardInterrupt:
                    print()  # 换行后继续
                    continue
                    
                if not user_input:
                    continue
                    
                # 添加到历史记录
                readline.add_history(user_input)
                
                # 检查退出命令
                if user_input.lower() in ['exit', 'quit']:
                    break
                    
                # 检查帮助命令
                if user_input.lower() in ['help', '--help', '-h']:
                    self.print_help()
                    continue
                
                # 处理记忆管理命令
                if user_input.lower().startswith('memory'):
                    command_parts = self.parse_command(user_input)
                    if self.handle_memory_command(command_parts):
                        continue
                
                # 处理 AI 命令
                if user_input.startswith('?'):
                    self.handle_ai_command(user_input)
                    continue
                
                # 执行普通命令
                command = self.parse_command(user_input)
                if command:
                    self.execute_command(command)
                    
            except Exception as e:
                print(f"错误: {e}")

def main():
    # 检查命令行参数
    if len(sys.argv) > 1 and sys.argv[1] in ['--help', '-h']:
        AIShell().print_help()
        return
        
    shell = AIShell()
    try:
        shell.run()
    finally:
        shell.save_history()

if __name__ == "__main__":
    main()
